# -*- coding: utf-8 -*-
# @Time    : 2020/8/19 下午9:03
# @Author  : lilong

"""
参考：https://blog.51cto.com/12597095/2497754
"""

import mitie


def get_vec():
    model_file = 'tmp_test/total_word_feature_extractor.dat'
    extractor = mitie.total_word_feature_extractor(model_file)
    feats = extractor.get_feature_vector("我")

    print(len(feats), len(feats))
    print("First 5 features of word '我'", feats)


def ner_train():
    from mitie import *

    sample = ner_training_instance(["I", "am", "looking", "for", "some", "cheap", "Mexican", "food", "."])

    sample.add_entity(xrange(5, 6), "pricerange")
    sample.add_entity(xrange(6, 7), "cuisine")

    sample2 = ner_training_instance(["show", "me", "indian", "restaurants", "in", "the", "centre", "."])
    sample2.add_entity(xrange(2, 3), "cuisine")
    sample2.add_entity(xrange(6, 7), "area")

    trainer = ner_trainer("/home/test/rasa_nlu_chi/data/total_word_feature_extractor_zh.dat")

    trainer.add(sample)
    trainer.add(sample2)

    trainer.num_threads = 4

    ner = trainer.train()

    ner.save_to_disk("new_ner_model.dat")

    tokens = ["I", "want", "expensive", "korean", "food"]
    entities = ner.extract_entities(tokens)

    for e in entities:
        range = e[0]
    tag = e[1]
    entity_text = " ".join(tokens[i] for i in range)
    print(" " + tag + ": " + entity_text)


if __name__ == '__main__':
    ner_train()
